import os
import json
from pdf2image import convert_from_path
import re
import pandas as pd
from tkinter import Tk, Label, Button, filedialog, messagebox
from tkinter.ttk import Progressbar
from paddleocr import PaddleOCR

# 加载行政区划代码的 JSON 文件
with open('regions.json', 'r', encoding='utf-8') as file:
    region_codes = json.load(file)

# 创建 PaddleOCR 实例
ocr = PaddleOCR(use_angle_cls=True, lang='ch')  # 使用简体中文

# 提取图片中的文本
def extract_text_from_image(image_path):
    result = ocr.ocr(image_path, cls=True)
    text = ""
    for line in result:
        for word_info in line:
            text += word_info[1][0] + "\n"  # 提取文本内容
    return text.strip()

# 从文本中提取发票信息
def extract_invoice_info(text, file_name):
    invoice_code_match = re.search(r'(\d{10})', text)
    invoice_code = invoice_code_match.group(1) if invoice_code_match else None

    # 从文件名中提取发票号码
    invoice_number = os.path.splitext(file_name)[0]  # 假设文件名即为发票号码

    # 使用更宽松的正则表达式匹配日期
    invoice_date_match = re.search(r'(\d{4})\s*[年./-]\s*(\d{1,2})\s*[月./-]\s*(\d{1,2})', text)
    invoice_date = (
        f"{invoice_date_match.group(1)}/{invoice_date_match.group(2).zfill(2)}/{invoice_date_match.group(3).zfill(2)}"
        if invoice_date_match else None
    )

    # 查找所有包含“￥”或“¥”符号的金额，并排序取第二大
    amounts = re.findall(r'[￥¥](\d+\.\d+)', text)
    print(amounts)
    print("================================================")
    amounts = sorted([float(amount) for amount in amounts], reverse=True)
    print(amounts)
    print("================================================")
    invoice_amount = amounts[1] if len(amounts) > 1 else (amounts[0] if amounts else None)

    invoice_type = '增值税专用发票'  # 发票类型固定

    # 检查发票代码是否符合要求：长度为10且前四位在region_codes中
    if invoice_code and len(invoice_code) == 10 and invoice_code[:4] in region_codes:
        region_name = region_codes[invoice_code[:4]]
    else:
        region_name = None  # 不符合条件时返回None

    return (
        invoice_code,
        invoice_number,
        invoice_date,
        invoice_type,
        invoice_amount,
        region_name
    )

# 处理 PDF，转换为图片并提取信息
def process_pdf(pdf_path, file_name):
    images = convert_from_path(pdf_path)
    for image in images:
        # 保存图片到临时文件，以便 PaddleOCR 读取
        temp_image_path = "temp_image.png"
        image.save(temp_image_path, "PNG")
        text = extract_text_from_image(temp_image_path)
        print(text)
        return extract_invoice_info(text, file_name)

# 遍历文件夹及其子文件夹中的所有PDF文件
def traverse_folders(folder_path, progress_bar):
    data = []
    pdf_files = []

    # 使用 os.walk() 递归遍历所有子文件夹
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if file.endswith('.pdf'):
                pdf_files.append(os.path.join(root, file))

    total_files = len(pdf_files)
    progress_step = 100 / total_files if total_files else 0

    # 处理每个 PDF 文件并更新进度条
    for i, pdf_path in enumerate(pdf_files):
        file_name = os.path.basename(pdf_path)
        invoice_data = process_pdf(pdf_path, file_name)
        data.append(invoice_data)
        progress_bar['value'] += progress_step
        progress_bar.update()
    return data

# 保存信息为 Excel 文件
def save_to_excel(data, folder_path):
    df = pd.DataFrame(data, columns=['发票代码', '发票号码', '开票日期yyyy/MM/dd', '发票类型', '发票金额（不含税）',
                                     '所属地区'])
    output_path = os.path.join(folder_path, '税票信息.xlsx')
    df.to_excel(output_path, index=False)
    return output_path

# 选择文件夹并开始处理
def select_folder():
    folder_path = filedialog.askdirectory()
    if not folder_path:
        return

    progress_bar['value'] = 0  # 重置进度条
    messagebox.showinfo("处理中", "开始提取发票信息，请稍等...")
    invoice_data = traverse_folders(folder_path, progress_bar)

    if not invoice_data:
        messagebox.showerror("错误", "未找到任何 PDF 文件！")
        return

    output_path = save_to_excel(invoice_data, folder_path)
    messagebox.showinfo("完成", f"税票信息已保存至 {output_path}！")

# 创建主界面
root = Tk()
root.title("发票信息提取器")
root.geometry("400x200")

Label(root, text="发票信息提取工具", font=("Arial", 16)).pack(pady=10)

Button(root, text="选择文件夹并处理", command=select_folder, font=("Arial", 12)).pack(pady=10)

progress_bar = Progressbar(root, orient='horizontal', length=300, mode='determinate')
progress_bar.pack(pady=20)

root.mainloop()
